package words;


import java.util.*;
import java.io.*;

import words.libsvm.intex_node;
import words.libsvm.svm;
import words.libsvm.svm_model;
import words.libsvm.svm_parameter;
import words.libsvm.svm_problem;

import words.libsvm.*;


public class IntexModel {

  String m_strTag;
  TagDefinition m_tagDef;

  svm_parameter m_params;
  svm_problem m_problem;
  svm_model m_model;


  public IntexModel(String tag_list, String strTag)
  {
    m_strTag = strTag;
    m_tagDef = new TagDefinition();
    m_tagDef.load(tag_list);
  }


  public boolean trainModel(String corpus_dir, String corpus_list)
  {
    readProblem(corpus_dir, corpus_list);
    setParams();

    String strError = svm.svm_check_parameter(m_problem, m_params);
    if(strError != null) {
      System.err.println("Error: " + strError +"!");
      return false;
    }

    m_model = svm.svm_train(m_problem, m_params);

    return true;
  }


  protected boolean readProblem(String corpus_dir, String corpus_list)
  {
    Vector vx = new Vector();
    Vector vy = new Vector();
    int[] counts = new int[2];

    try {
      BufferedReader br = new BufferedReader(new FileReader(corpus_list));
      String strLine = br.readLine();
      while (strLine != null) {
        strLine = strLine.trim();
        if (strLine.length() > 0) {
          // load input document
          String strInFile = corpus_dir + "/" + strLine;
          
          System.out.print("processing [" + strInFile + "] ... ");
	  
          SgmlDocument sgmlDoc = new SgmlDocument(m_tagDef);
          sgmlDoc.load(strInFile);
          SgmlTokSentDocument sgmlTSDoc = new SgmlTokSentDocument(sgmlDoc);
          // sgmlTSDoc.loadPOSTags(strInFile + ".pos");
          
          // create protein interaction document
          IntexDocument intexDoc = new IntexDocument(sgmlTSDoc, m_strTag);
	  // store all interaction examples in vx, vy
	  updateProblem(intexDoc, vx, vy);
	  // update counts
	  int[] dCounts = intexDoc.pairsCounts();
	  for (int i = 0; i < counts.length; i++)
	    counts[i] += dCounts[i];

	  System.out.println("done!");
        }	  
        strLine = br.readLine();
      }
    }
    catch (IOException e) {
      System.err.println(e);
      return false;
    }

    m_problem = new svm_problem();
    m_problem.l = vy.size();
    m_problem.x = (intex_node[][]) vx.toArray(new intex_node[0][]);
    m_problem.y = new double[vy.size()];
    for (int i = 0; i < vy.size(); i++)
      m_problem.y[i] = ((Double) vy.get(i)).doubleValue();

    // print counts
    for (int i = 0; i < counts.length; i++)
      System.out.println(counts[i] + " pairs with label " + i);

    return true;
  }


  protected boolean updateProblem(IntexDocument intexDoc, Vector vx, Vector vy)
  {
    for (int i = 0; i < intexDoc.m_sents.length; i++) {
      IntexSentence sent = intexDoc.m_sents[i];
      // read upper-triangular matrix of entity pairs
      for (int j = 0; j < sent.m_entities.length; j++) {
	Annotation anno1 = sent.m_entities[j];
	for (int k = j + 1; k < sent.m_entities.length; k++) {
	  Annotation anno2 = sent.m_entities[k];
	  if (sent.m_pairs[j][k] != null) {
	    // store class information
	    vy.add(new Double(sent.m_pairs[j][k].m_nTrueLabel));

	    // create relation example sentence
	    assert anno1.m_nBegin + anno1.m_nLen < anno2.m_nBegin + anno2.m_nLen; 
	    Vector vecSent = new Vector();
	    for (int w = 0; w < anno1.m_nBegin; w++)
	      vecSent.add(sent.m_sgmlTS.m_vecTokens[w]);
	    vecSent.add("<p>");
	    for (int w = anno1.m_nBegin + anno1.m_nLen; w < anno2.m_nBegin; w++)
	      vecSent.add(sent.m_sgmlTS.m_vecTokens[w]);
	    vecSent.add("<p>");
	    for (int w = anno2.m_nBegin + anno2.m_nLen; 
		 w < sent.m_sgmlTS.m_vecTokens.length; w++)
	      vecSent.add(sent.m_sgmlTS.m_vecTokens[w]);

	    intex_node[] inode = new intex_node[1];
	    inode[0] = new intex_node();
	    inode[0].m_value = (String[]) vecSent.toArray(new String[0]);
	    vx.add(inode);
	  }
	}
      }
    }

    return true;
  }


  protected boolean setParams()
  {
    m_params = new svm_parameter();

    // default values
    m_params.svm_type = svm_parameter.C_SVC;
    m_params.kernel_type = svm_parameter.RELATION;
    m_params.degree = 3;
    m_params.gamma = 0;
    m_params.coef0 = 0;
    m_params.nu = 0.5;
    m_params.cache_size = 40;
    m_params.C = 1;
    m_params.eps = 1e-3;
    m_params.p = 0.1;
    m_params.shrinking = 1;
    m_params.probability = 1;
    m_params.nr_weight = 0;
    m_params.weight_label = new int[0];
    m_params.weight = new double[0];

    return true;
  }


  public boolean testModel(String corpus_dir, String corpus_list, String out_dir)
  {
    try {
      BufferedReader br = new BufferedReader(new FileReader(corpus_list));
      String strLine = br.readLine();
      while (strLine != null) {
        strLine = strLine.trim();
        if (strLine.length() > 0) {
          // load input document
          String strInFile = corpus_dir + "/" + strLine;
          
          System.out.print("processing [" + strInFile + "] ... ");

          SgmlDocument sgmlDoc = new SgmlDocument(m_tagDef);
          sgmlDoc.load(strInFile);
          SgmlTokSentDocument sgmlTSDoc = new SgmlTokSentDocument(sgmlDoc);
          // sgmlTSDoc.loadPOSTags(strInFile + ".pos");
          
          // create protein interaction document
          IntexDocument intexDoc = new IntexDocument(sgmlTSDoc, m_strTag);
	  inferRelations(intexDoc);

	  // write output document
	  String strOutFile = out_dir + "/" + strLine + ".out";
	  intexDoc.writeInf(strOutFile);

	  System.out.println("done!");
        }	  
        strLine = br.readLine();
      }
    }
    catch (IOException e) {
      System.err.println(e);
      return false;
    }
    
    return true;
  }


  protected boolean inferRelations(IntexDocument intexDoc)
  {
    // label probabilities
    double[] probs = new double[2];

    for (int i = 0; i < intexDoc.m_sents.length; i++) {
      IntexSentence sent = intexDoc.m_sents[i];
      // read upper-triangular matrix of entity pairs
      for (int j = 0; j < sent.m_entities.length; j++) {
	Annotation anno1 = sent.m_entities[j];
	for (int k = j + 1; k < sent.m_entities.length; k++) {
	  Annotation anno2 = sent.m_entities[k];
	  if (sent.m_pairs[j][k] != null) {
	    // create relation example sentence
	    assert anno1.m_nBegin + anno1.m_nLen < anno2.m_nBegin + anno2.m_nLen; 
	    Vector vecSent = new Vector();
	    for (int w = 0; w < anno1.m_nBegin; w++)
	      vecSent.add(sent.m_sgmlTS.m_vecTokens[w]);
	    vecSent.add("<p>");
	    for (int w = anno1.m_nBegin + anno1.m_nLen; w < anno2.m_nBegin; w++)
	      vecSent.add(sent.m_sgmlTS.m_vecTokens[w]);
	    vecSent.add("<p>");
	    for (int w = anno2.m_nBegin + anno2.m_nLen; 
		 w < sent.m_sgmlTS.m_vecTokens.length; w++)
	      vecSent.add(sent.m_sgmlTS.m_vecTokens[w]);

	    intex_node[] inode = new intex_node[1];
	    inode[0] = new intex_node();
	    inode[0].m_value = (String[]) vecSent.toArray(new String[0]);
	    double dd = svm.svm_predict_probability(m_model, inode , probs);
	    int nLabel  = (int) dd;
	    int nr_class = svm.svm_get_nr_class(m_model);
	    int[] labels = new int[nr_class];
	    svm.svm_get_labels(m_model, labels);
	    for (int l = 0; l < nr_class; l++)
	      if (labels[l] == nLabel) {
		sent.m_pairs[j][k].m_conf = probs[l];
		break;
	      }
	    sent.m_pairs[j][k].m_nInfLabel = nLabel;
	  }
	}
      }
    }

    return true;    
  }


  static public void main_train(String[] args)
  {
    if (args.length != 4) {
      System.out.println("Usage: java IntexModel <corpus-dir> <corpus-list> " + 
                         "<tag-list> <intex-model>!");
      System.exit(1);
    }
    
    IntexModel intex = new IntexModel(args[2], "prot");
    intex.trainModel(args[0], args[1]);

    try {
      svm.svm_save_model(args[3], intex.m_model);    
    }
    catch (IOException e) {
      System.err.println(e);
    }
  }
  

  static public void main_test(String[] args)
  {
    if (args.length != 5) {
      System.out.println("Usage: java IntexModel <corpus-dir> <corpus-list> " + 
                         "<tag-list> <svm-model> <out-dir>!");
      System.exit(1);
    }

    IntexModel intex = new IntexModel(args[2], "prot");
    try {
      intex.m_model = svm.svm_load_model(args[3]);
    }
    catch (IOException e) {
      System.err.println(e);
    }

    intex.testModel(args[0], args[1], args[4]);
  }


  static public void main(String[] args)
  {
    if (args.length == 0) {
      System.out.println("Usage: java IntexModel -[train|test] ...");
      System.exit(1);
    }
    
    String[] newArgs = new String[args.length - 1];
    for (int i = 1; i < args.length; i++)
      newArgs[i-1] = new String(args[i]);
    
    if (args[0].compareTo("-train") == 0) {
      main_train(newArgs);
      return;
    }
    if (args[0].compareTo("-test") == 0) {
      main_test(newArgs);
      return;
    }
  }


}
